from ultralytics import YOLO
from cnocr import CnOcr
import re
import torch
import cv2
from PIL import Image, ImageTk
import numpy as np
import tkinter as tk
from tkinter import filedialog, ttk
import threading
import os

# 初始化模型
device = 'cuda' if torch.cuda.is_available() else 'cpu'


class LicensePlateRecognitionApp:
    def __init__(self, root):
        self.root = root
        self.root.title("车牌识别系统")
        self.root.geometry("800x600")
        self.root.configure(bg="white")

        # 顶部标题
        title_frame = ttk.Frame(root)
        title_frame.pack(fill=tk.X, padx=10, pady=5)

        title_label = ttk.Label(
            title_frame,
            text="车牌识别系统",
            foreground="black",
            font=("Microsoft YaHei", 16, "bold")
        )
        title_label.pack(pady=10)

        # 原始图像显示区
        image_frame = ttk.Frame(root)
        image_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=5)

        # 创建画布显示原始图像
        image_label = ttk.Label(
            image_frame,
            text="原始图像",
            foreground="black",
            font=("Microsoft YaHei", 11)
        )
        image_label.pack(anchor=tk.NW)

        self.canvas_original = tk.Canvas(
            image_frame,
            bg="white",
            highlightthickness=0,
            height=300
        )
        self.canvas_original.pack(fill=tk.BOTH, expand=True)
        self.original_img = None
        self.original_img_tk = None

        # 控制按钮区域
        button_frame = ttk.Frame(root)
        button_frame.pack(fill=tk.X, padx=20, pady=10)

        # 选择图片按钮
        self.btn_select = ttk.Button(
            button_frame,
            text="选择图片",
            command=self.select_image,
            width=10
        )
        self.btn_select.pack(side=tk.LEFT, padx=(0, 15))

        # 开始识别按钮 - 蓝色背景，黑色字体
        self.btn_start = ttk.Button(
            button_frame,
            text="开始识别",
            command=self.start_recognition,
            width=10
        )
        self.btn_start.pack(side=tk.LEFT)

        # 识别结果显示区域
        result_frame = ttk.LabelFrame(root, text="识别结果")
        result_frame.pack(fill=tk.X, padx=20, pady=10)

        self.lbl_results = ttk.Label(
            result_frame,
            text="浙FM3568",  # 默认示例车牌
            foreground="black",
            background="white",
            font=("Microsoft YaHei", 24, "bold"),
            anchor=tk.CENTER
        )
        self.lbl_results.pack(fill=tk.X, pady=10)

        # 状态信息区域
        status_frame = ttk.Frame(root)
        status_frame.pack(fill=tk.X, side=tk.BOTTOM, padx=20, pady=10)

        # 文件路径信息
        self.path_var = tk.StringVar()
        self.path_var.set("已选择图像: test.png")
        path_label = ttk.Label(
            status_frame,
            textvariable=self.path_var,
            foreground="black",
            font=("Microsoft YaHei", 9)
        )
        path_label.pack(side=tk.LEFT)

        # 状态信息
        self.status_var = tk.StringVar()
        self.status_var.set("状态: 准备就绪")
        status_label = ttk.Label(
            status_frame,
            textvariable=self.status_var,
            foreground="black",
            font=("Microsoft YaHei", 9)
        )
        status_label.pack(side=tk.RIGHT)

        # 初始化变量
        self.is_processing = False
        self.models_loaded = False

        # 加载默认图片
        self.load_default_image()

    def load_default_image(self):
        """加载默认图片"""
        self.img_path = "test.png"
        self.display_image(self.img_path)
        self.lbl_results.config(text="浙FM3568")  # 显示示例车牌
        self.status_var.set("状态: 准备就绪")

    def display_image(self, img_path):
        """显示图像到原始图像区域"""
        try:
            # 加载并显示原始图像
            self.original_img = cv2.imread(img_path)
            if self.original_img is None:
                raise ValueError(f"无法加载图像: {img_path}")

            # 转换颜色通道
            self.original_img = cv2.cvtColor(self.original_img, cv2.COLOR_BGR2RGB)

            # 将图片调整到合适尺寸
            img_pil = Image.fromarray(self.original_img)
            img_pil.thumbnail((600, 300), Image.LANCZOS)  # 限制在画布范围内

            # 在画布上显示图片
            self.original_img_tk = ImageTk.PhotoImage(img_pil)
            self.canvas_original.delete("all")
            self.canvas_original.create_image(
                (self.canvas_original.winfo_width() - img_pil.width) // 2,
                0,
                anchor=tk.N,
                image=self.original_img_tk
            )

            # 更新状态
            filename = os.path.basename(img_path)
            self.path_var.set(f"已选择图像: {filename}")

        except Exception as e:
            self.status_var.set(f"状态: 图片加载错误 - {str(e)}")

    def select_image(self):
        """选择图片文件"""
        if self.is_processing:
            self.status_var.set("状态: 正在处理中，请稍后")
            return

        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp")]
        )
        if file_path:
            self.img_path = file_path
            self.display_image(file_path)
            self.lbl_results.config(text="等待识别结果...")  # 清除之前的识别结果
            self.status_var.set("状态: 图片已加载")

    def start_recognition(self):
        """开始识别过程"""
        if not hasattr(self, 'img_path') or not self.img_path:
            self.status_var.set("状态: 请先选择图片")
            return

        if self.is_processing:
            self.status_var.set("状态: 正在处理中")
            return

        self.status_var.set("状态: 开始识别中...")
        self.lbl_results.config(text="识别中...")
        self.is_processing = True

        # 在新线程中执行识别
        threading.Thread(target=self.recognize, daemon=True).start()

    def recognize(self):
        """执行车牌识别"""
        try:
            # 检查是否需要加载模型
            if not self.models_loaded:
                self.load_models()

            # 使用YOLO检测车牌位置
            img = cv2.cvtColor(self.original_img, cv2.COLOR_RGB2BGR)
            results = self.yolo_model.predict(img, verbose=False)

            plate_text = None
            for result in results:
                boxes = result.boxes
                for box in boxes:
                    if float(box.conf) >= 0.5:  # 置信度阈值
                        # 获取坐标并转为整数
                        x1, y1, x2, y2 = map(int, box.data[0][:4])

                        # 裁剪车牌区域
                        cropped_plate = img[y1:y2, x1:x2]

                        # OCR识别
                        plate_text, _ = self.recognize_plate_from_crop(cropped_plate)
                        if plate_text:
                            break  # 使用第一个识别出的车牌

                if plate_text:
                    break

            # 更新UI
            if plate_text:
                self.root.after(10, lambda: self.lbl_results.config(text=plate_text))
                self.root.after(10, lambda: self.status_var.set("状态: 识别完成"))
            else:
                self.root.after(10, lambda: self.lbl_results.config(text="未识别到车牌"))
                self.root.after(10, lambda: self.status_var.set("状态: 识别完成，未找到车牌"))

        except Exception as e:
            self.root.after(10, lambda: self.status_var.set(f"状态: 识别错误 - {str(e)}"))
            self.root.after(10, lambda: self.lbl_results.config(text="识别失败"))
        finally:
            self.is_processing = False
            self.root.after(10, lambda: self.status_var.set("状态: 就绪"))

    def load_models(self):
        """加载YOLO和OCR模型"""
        try:
            self.status_var.set("状态: 模型加载中...")
            self.yolo_model = YOLO('yolov8n.pt')
            self.ocr = CnOcr(
                rec_model_name='densenet_lite_136-gru',
                det_model_name='db_resnet34',
                cand_alphabet="0123456789ABCDEFGHJKLMNPQRSTUVWXYZ京沪津渝冀晋辽吉黑苏浙皖闽赣鲁豫鄂湘粤琼川贵云陕甘青蒙桂宁新藏使领",
                context='gpu' if device == 'cuda' else 'cpu'
            )
            self.models_loaded = True
            self.status_var.set("状态: 模型加载完成")
        except Exception as e:
            self.status_var.set(f"状态: 模型加载失败 - {str(e)}")

    def clean_license_plate_text(self, text):
        """清理车牌文本，只保留中文、大写字母和数字"""
        if not text:
            return ""
        return re.sub(r'[^\u4e00-\u9fa5A-Z0-9]', '', text.upper())

    def recognize_plate_from_crop(self, cropped_img):
        """识别裁剪后的车牌图像"""
        try:
            result = self.ocr.ocr(cropped_img)
            if not result:
                return None, None

            # 获取最匹配的结果
            sorted_results = sorted(result, key=lambda x: x.get('score', 0), reverse=True)
            if sorted_results:
                best_result = sorted_results[0]
                cleaned_text = self.clean_license_plate_text(best_result['text'])
                return cleaned_text, best_result.get('score', 0)

            return None, None

        except Exception as e:
            self.status_var.set(f"OCR错误: {str(e)}")
            return None, None


# 运行应用程序
if __name__ == "__main__":
    root = tk.Tk()
    app = LicensePlateRecognitionApp(root)
    root.mainloop()